package com.hlx.config;

import com.hlx.constants.SystemPrompt;
import com.hlx.tools.MajorTool;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class ChatConfiguration {

	/**
	 * 由于有ollamaEmbeddingModel和openAiEmbeddingModel两个bean，
	 * 所以这里将OpenAiEmbeddingModel的bean设置为默认bean
	 * @param model
	 * @return
	 */
	@Bean
	public EmbeddingModel embeddingModel(OpenAiEmbeddingModel model) {
		return model;
	}

	/**
	 * AI聊天的记忆对象
	 * @param chatMemoryRepository
	 * @return
	 */
	@Bean
	public ChatMemory chatMemory(JdbcChatMemoryRepository chatMemoryRepository) {
		return MessageWindowChatMemory.builder()
				.chatMemoryRepository(chatMemoryRepository)
				.maxMessages(20)
				.build();
	}

	/**
	 * 哄哄模拟器的记忆对象
	 * @return
	 */
	@Bean
	public ChatMemory gameChatMemory() {
		return MessageWindowChatMemory.builder()
				.maxMessages(20)
				.build();
	}

	/**
	 * 创建AI聊天的ChatClient实例
	 * @param model
	 * @param chatMemory
	 * @return
	 */
	@Bean
	public ChatClient chatClient(OllamaChatModel model, ChatMemory chatMemory) {
		return ChatClient
				.builder(model) // 创建一个 ChatClient工厂实例
				.defaultSystem(SystemPrompt.SYSTEM_CHAT_PROMPT) // 角色预设
				.defaultAdvisors(new SimpleLoggerAdvisor(),
						MessageChatMemoryAdvisor.builder(chatMemory).build()) // 配置日志记录器advisor
				.build(); // 构建ChatClient实例
	}

	/**
	 * 创建哄哄模拟器的ChatClient实例
	 * @param model
	 * @param gameChatMemory
	 * @return
	 */
	@Bean
	public ChatClient gameChatClient(OpenAiChatModel model, ChatMemory gameChatMemory){
		return ChatClient
				.builder(model)
				.defaultSystem(SystemPrompt.SYSTEM_GAME_PROMPT)
				.defaultAdvisors(new SimpleLoggerAdvisor(),
						MessageChatMemoryAdvisor.builder(gameChatMemory).build())
				.build();
	}

	/**
	 * 创建智能客服的ChatClient实例
	 * @param model
	 * @param chatMemory
	 * @param majorTool
	 * @return
	 */
	@Bean
	public ChatClient serviceChatClient(OpenAiChatModel model, ChatMemory chatMemory, MajorTool majorTool){
		return ChatClient
				.builder(model)
				.defaultSystem(SystemPrompt.SYSTEM_SERVICE_PROMPT)
				.defaultAdvisors(new SimpleLoggerAdvisor(),
						MessageChatMemoryAdvisor.builder(chatMemory).build())
				.defaultTools(majorTool)
				.build();
	}

	/**
	 * 创建PDF阅读器的ChatClient实例
	 * @param model
	 * @param chatMemory
	 * @param vectorStore
	 * @return
	 */
	@Bean
	public ChatClient pdfChatClient(OpenAiChatModel model, ChatMemory chatMemory, VectorStore vectorStore){
		return ChatClient.builder(model)
				.defaultSystem(SystemPrompt.SYSTEM_PDF_PROMPT)
				.defaultAdvisors(
						new SimpleLoggerAdvisor(),
						MessageChatMemoryAdvisor.builder(chatMemory).build(),
						QuestionAnswerAdvisor.builder(vectorStore)
								.searchRequest(SearchRequest.builder()
										.similarityThreshold(0.6)
										.topK(2)
										.build())
								.build()
				)
				.build();
	}
}
